# -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import print_function
import matplotlib.pyplot as plt
from matplotlib.pyplot import MultipleLocator  # 从pyplot导入MultipleLocator类，这个类用于设置刻度间隔
from matplotlib.patches import Circle
import matplotlib.patches as mpathes
import matplotlib.patches as patches
import matplotlib.tri as tri
import csv
import os
import sys
import math
from six.moves import range
from six.moves import zip

sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__))))
dirPath = os.path.dirname(os.path.realpath(__file__)) # /home/chiweiming/catkin_ws/src/cwm_multi_nav/cwm_nav/script

class Processor(object):
    def __init__(self):
        self.key1 = 0
        self.all_total_time = 0.0
        self.all_auto_ctrl_time = 0.0
        self.all_bci_ctrl_time = 0.0
        self.total_csv_files = 0
        self.total_mean_leader_lin_vel = 0.0
        self.total_MPC_ensure_safety_time = 0.0

    # 绘制每个曲线，输出相关计算参数
    def process_all(self, path):
        for root,dirs,files in os.walk(path):
            files.sort()
            total_files = len(files)
            # print(files)
            # print("total_files: %d" %total_files)
            for file in files:
                if (file[-3:] == 'csv'):
                    self.total_csv_files = self.total_csv_files + 1
                    print(("process: %s" %file))
                    self.plot_traj(cur_file_path=path + "/" + file)
            
            # 指标总平均
            total_mean_time = 0
            mean_lin_vel = 0.0
            total_MPC_safety_time_mean = 0.0
            if (self.total_csv_files != 0):
                total_mean_time = self.all_total_time / self.total_csv_files
                mean_lin_vel = self.total_mean_leader_lin_vel / self.total_csv_files
                total_MPC_safety_time_mean = self.total_MPC_ensure_safety_time / self.total_csv_files
            total_MPC_safety_prop_mean = total_MPC_safety_time_mean / self.all_total_time
            self.all_auto_ctrl_prop = self.all_auto_ctrl_time / self.all_total_time
            self.all_bci_ctrl_prop = self.all_bci_ctrl_time / self.all_total_time
            
            print("")
            print(("total run time: %lf" %self.all_total_time))
            print(("mean run time: %lf" %total_mean_time))
            print(("mean lin vel: %lf" %mean_lin_vel))
            print(("all_auto_ctrl_time: %lf, all_auto_ctrl_prop: %lf" %(self.all_auto_ctrl_time, self.all_auto_ctrl_prop)))
            print(("all_bci_ctrl_time: %lf, all_bci_ctrl_prop: %lf" %(self.all_bci_ctrl_time, self.all_bci_ctrl_prop)))
            print(("MPC_ensure_safety_time: %lf, MPC_ensure_safety_time_prop: %lf" %(total_MPC_safety_time_mean, total_MPC_safety_prop_mean)))
            
            real_auto_time = self.all_auto_ctrl_time + self.total_MPC_ensure_safety_time
            real_auto_prop = real_auto_time / self.all_total_time
            real_bci_time = self.all_bci_ctrl_time - self.total_MPC_ensure_safety_time
            real_bci_prop = real_bci_time / self.all_total_time
            print(("real Auto: %lf, prop: %lf"%(real_auto_time, real_auto_prop)))
            print(("real bci: %lf, prop: %lf"%(real_bci_time, real_bci_prop)))


    def plot_traj(self, cur_file_path):
        leader_x = []
        leader_y = []
        leader_v = []
        follower1_x = []
        follower1_y = []
        follower2_x = []
        follower2_y = []

        auto_flag_list = []
        leader_ctrl_lin_vel_input = []
        leader_ctrl_ang_vel_input = []
        leader_ctrl_lin_vel_output = []
        leader_ctrl_ang_vel_output = []

        with open(cur_file_path, "r") as csvfile:
            reader = csv.reader(csvfile)
            # 这里不需要readlines
            for line in reader:
                leader_y.append(float(line[0]))
                leader_x.append(-1.0 * float(line[1])) 
                leader_v.append(float(line[2]))
                follower1_y.append(float(line[4]))
                follower1_x.append(-1.0 * float(line[5])) 
                follower2_y.append(float(line[8])) 
                follower2_x.append(-1.0 * float(line[9]))
                auto_flag_list.append(float(line[13]))
                leader_ctrl_lin_vel_input.append(float(line[14]))
                leader_ctrl_ang_vel_input.append(float(line[15]))
                leader_ctrl_lin_vel_output.append(float(line[16]))
                leader_ctrl_ang_vel_output.append(float(line[17]))

        plt.plot(leader_x, leader_y, linewidth=2.0, color='black', label="$trajectory1$")
        plt.plot(follower1_x, follower1_y, linewidth=0.8, color='black')     # 跟随者1轨迹
        plt.plot(follower2_x, follower2_y, linewidth=0.8, color='black')     # 跟随者2轨迹

        # 绘制编队队形
        for l1_x, l1_y, f1_x, f1_y, f2_x, f2_y in zip(leader_x, leader_y, follower1_x, follower1_y, follower2_x, follower2_y):
            if ((self.key1 % 50) == 0) or (self.key1 == len(leader_x) -1) :
                X = []
                Y = []
                X.append(l1_x)
                X.append(f1_x)
                X.append(f2_x)
                Y.append(l1_y)
                Y.append(f1_y)
                Y.append(f2_y)
                triangles = tri.Triangulation(X, Y)
                plt.triplot(triangles, color='gray', linestyle=(0, (5, 3)), linewidth=1.0)
            self.key1 = self.key1 + 1

        # 设置目标点和障碍物
        # 目标点
        # plt.legend(frameon=True, loc="upper right", fontsize='small')  # 分别为图例无边框、图例放在右上角、图例大小
        # x_major_locator = MultipleLocator(1)                    # 设置显示的间距
        # y_major_locator = MultipleLocator(1)
        ax = plt.gca()                                          # ax为两条坐标轴的实例
        # ax.xaxis.set_major_locator(x_major_locator)             # 把x轴的主刻度设置为1的倍数
        # ax.yaxis.set_major_locator(y_major_locator)             # 把y轴的主刻度设置为1的倍数
        plt.axis('equal')                                       # 把坐标轴单位长度都变的一样
        plt.xlim(-8.0, 8.0)  # 把x轴的刻度范围设置为-7.5到7.5，因为0.5不满一个刻度间隔，所以数字不会显示出来，但是能看到一点空白
        plt.ylim(-8.0, 7.0)
        ax.add_artist(Circle(xy=(-5.0, 5.0), radius=0.4, color='orangered'))
        ax.add_artist(Circle(xy=(0.0, 5.0), radius=0.4, color='orangered'))
        ax.add_artist(Circle(xy=(5.0, 5.0), radius=0.4, color='orangered'))

        # 障碍物
        ax.add_artist(Circle(xy=(1.0, -5), radius=0.2, color='g'))   # A
        ax.add_artist(Circle(xy=(-3, -6), radius=0.2, color='g'))  # B
        ax.add_artist(Circle(xy=(4, -3), radius=0.2, color='g'))   # C
        ax.add_artist(Circle(xy=(-6, -2), radius=0.2, color='g'))  # D
        ax.add_artist(Circle(xy=(7, -1), radius=0.2, color='g'))   # E
        ax.add_artist(Circle(xy=(0, -1), radius=0.2, color='g'))   # F
        ax.add_artist(Circle(xy=(-3, -1), radius=0.2, color='g'))  # G
        ax.add_artist(Circle(xy=(-8, 0), radius=0.2, color='g'))   # H
        ax.add_artist(Circle(xy=(4, 1), radius=0.2, color='g'))    # I
        ax.add_artist(Circle(xy=(-5, 1), radius=0.2, color='g'))   # J
        ax.add_artist(Circle(xy=(0, 3), radius=0.2, color='g'))    # K
        ax.add_artist(Circle(xy=(-3, 3), radius=0.2, color='g'))   # L
        ax.add_artist(Circle(xy=(-2, 1), radius=0.2, color='g'))   # X


        # print(cur_file_path[:-3]+ "png") 
        # plt.savefig(cur_file_path[:-3] + "png", dpi=600)  # 保存图片
        plt.clf()
        # plt.show()

        # 平均速度
        self.mean_lin_vel(leader_v)

        # 控制占比
        self.control_prop(auto_flag_list, leader_ctrl_lin_vel_input, leader_ctrl_ang_vel_input,\
            leader_ctrl_lin_vel_output, leader_ctrl_ang_vel_output)
    
    def mean_lin_vel(self, leader_v):      
        len_v = len(leader_v)
        sum_leader_v = 0.0
        for i in range(len_v):
            sum_leader_v += leader_v[i]
        mean_v = sum_leader_v / len_v
        self.total_mean_leader_lin_vel = self.total_mean_leader_lin_vel + mean_v
        print(("    Mean linear velocity: %lf" %mean_v))

    def control_prop(self, auto_flag_list, \
                    leader_ctrl_lin_vel_input, leader_ctrl_ang_vel_input, \
                    leader_ctrl_lin_vel_output, leader_ctrl_ang_vel_output):
            # print(len(auto_flag_list))
            auto_ctrl_count = 0
            bci_ctrl_count = 0
            MPC_ensure_safety_count = 0
            MPC_ensure_safety_count_v = 0
            MPC_ensure_safety_count_w = 0
            for i in range(len(auto_flag_list)):
                # auto_flag_list[i] == 1 自主控制状态
                if (abs(auto_flag_list[i] - 1.0) <= 0.05):
                    auto_ctrl_count = auto_ctrl_count + 1
                elif (abs(auto_flag_list[i]) <= 0.05):
                    bci_ctrl_count = bci_ctrl_count + 1
                    if (abs(leader_ctrl_lin_vel_input[i] - leader_ctrl_lin_vel_output[i]) >= 0.005):
                        MPC_ensure_safety_count_v += 1
                    if (abs(leader_ctrl_ang_vel_input[i] - leader_ctrl_ang_vel_output[i]) >= 0.005):
                        MPC_ensure_safety_count_w += 1
                    if (abs(leader_ctrl_lin_vel_input[i] - leader_ctrl_lin_vel_output[i]) >= 0.005) or\
                        (abs(leader_ctrl_ang_vel_input[i] - leader_ctrl_ang_vel_output[i]) >= 0.005):
                        MPC_ensure_safety_count += 1

            total_time = len(auto_flag_list) / 10.0
            auto_ctrl_time = auto_ctrl_count / 10.0
            auto_ctrl_prop = auto_ctrl_time / total_time
            bci_ctrl_time = bci_ctrl_count / 10.0
            bci_ctrl_prop = bci_ctrl_time / total_time
            MPC_ensure_safety_time_v = MPC_ensure_safety_count_v / 10.0
            MPC_ensure_safety_time_w = MPC_ensure_safety_count_w / 10.0
            MPC_ensure_safety_time = MPC_ensure_safety_count / 10.0
            MPC_ensure_safety_time_prop = MPC_ensure_safety_time / total_time
            self.total_MPC_ensure_safety_time += MPC_ensure_safety_time
            self.all_total_time = self.all_total_time + total_time
            self.all_auto_ctrl_time = self.all_auto_ctrl_time + auto_ctrl_time
            self.all_bci_ctrl_time = self.all_bci_ctrl_time + bci_ctrl_time

            print(("    run time: %lf" %total_time))
            print(("    auto_ctrl_time: %lf, auto_ctrl_prop: %lf" %(auto_ctrl_time, auto_ctrl_prop)))
            print(("    bci_ctrl_time: %lf, bci_ctrl_prop: %lf" %(bci_ctrl_time, bci_ctrl_prop)))

            print(("    MPC_ensure_safety_time: %lf, MPC_ensure_safety_time_prop: %lf" %(MPC_ensure_safety_time, MPC_ensure_safety_time_prop)))
            print(("    MPC_ensure_safety_time_v: %lf, MPC_ensure_safety_time_w: %lf" %(MPC_ensure_safety_time_v, MPC_ensure_safety_time_w)))

    # 把多个轨迹绘制到一张图上
    def plot_trajs(self, path, auto):
        # bci
        # traj1_path = path + "/Hands-A-1.csv"
        # traj2_path = path + "/Hands-B-2.csv"
        # traj3_path = path + "/Hands-C-1.csv"

        # # auto
        # traj1_path = path + "/Hands-A-1.csv"
        # traj2_path = path + "/Hands-B-3.csv"
        # traj3_path = path + "/Hands-C-3.csv"
        traj1_path = path + "/yzg-brain_auto_A-1.csv"
        traj2_path = path + "/yzg-brain_auto_A-1.csv"
        traj3_path = path + "/yzg-brain_auto_A-1.csv"
        self.plot_one_traj(traj1_path, traj_color='black', formation_color='gray', auto=auto)
        self.plot_one_traj(traj2_path, traj_color='black', formation_color='gray', auto=auto)
        self.plot_one_traj(traj3_path, traj_color='black', formation_color='gray', auto=auto)

        self.plot_env()

        # print(cur_file_path[:-3]+ "png") 
        plt.savefig(path + "/multi_trajs.png", dpi=600)  # 保存图片
        plt.show()
        plt.clf()
    
    def plot_one_traj(self, path, traj_color, formation_color, auto):
        leader_x = []
        leader_y = []
        follower1_x = []
        follower1_y = []
        follower2_x = []
        follower2_y = []
        auto_flag_list = []

        with open(path, "r") as csvfile:
            reader = csv.reader(csvfile)
            next(reader)
            # 这里不需要readlines
            for line in reader:
                leader_y.append(float(line[0]))
                leader_x.append(-1.0 * float(line[1])) 
                follower1_y.append(float(line[4]))
                follower1_x.append(-1.0 * float(line[5])) 
                follower2_y.append(float(line[8])) 
                follower2_x.append(-1.0 * float(line[9]))
                auto_flag_list.append(float(line[13]))
        if (auto):
            plot_traj_x = []
            plot_traj_y = []
            plot_traj_f1_x = []
            plot_traj_f1_y = []
            plot_traj_f2_x = []
            plot_traj_f2_y = []
            for i in range(1, len(auto_flag_list)):
                plot_traj_x.append(leader_x[i])
                plot_traj_y.append(leader_y[i])
                plot_traj_f1_x.append(follower1_x[i])
                plot_traj_f1_y.append(follower1_y[i])
                plot_traj_f2_x.append(follower2_x[i])
                plot_traj_f2_y.append(follower2_y[i])
                if ((abs(auto_flag_list[i] - auto_flag_list[i - 1]) >= 0.3) or (i ==len(auto_flag_list)-1)):
                    if ((i!= len(auto_flag_list)-1) and (abs(auto_flag_list[i] - 1.0) <= 0.05)):
                        plt.plot(plot_traj_x, plot_traj_y, linewidth=2.0, color=traj_color, label="$trajectory1$")
                        plt.plot(plot_traj_f1_x, plot_traj_f1_y, linewidth=0.8, color=traj_color)     # 跟随者1轨迹
                        plt.plot(plot_traj_f2_x, plot_traj_f2_y, linewidth=0.8, color=traj_color)     # 跟随者2轨迹
                    elif ((i ==len(auto_flag_list)-1) and auto_flag_list[i] <= 0.05):
                        plt.plot(plot_traj_x, plot_traj_y, linewidth=2.0, color=traj_color, label="$trajectory1$")
                        plt.plot(plot_traj_f1_x, plot_traj_f1_y, linewidth=0.8, color=traj_color)     # 跟随者1轨迹
                        plt.plot(plot_traj_f2_x, plot_traj_f2_y, linewidth=0.8, color=traj_color)     # 跟随者2轨迹
                    elif ((i!= len(auto_flag_list)-1) and (auto_flag_list[i] <= 0.05)):
                        plt.plot(plot_traj_x, plot_traj_y, linewidth=2.0, color="red", label="$trajectory1$")
                        plt.plot(plot_traj_f1_x, plot_traj_f1_y, linewidth=0.8, color=traj_color)     # 跟随者1轨迹
                        plt.plot(plot_traj_f2_x, plot_traj_f2_y, linewidth=0.8, color=traj_color)     # 跟随者2轨迹
                    elif (i ==len(auto_flag_list)-1) and (abs(auto_flag_list[i] - 1.0) <= 0.05):
                        plt.plot(plot_traj_x, plot_traj_y, linewidth=2.0, color="red", label="$trajectory1$")
                        plt.plot(plot_traj_f1_x, plot_traj_f1_y, linewidth=0.8, color=traj_color)     # 跟随者1轨迹
                        plt.plot(plot_traj_f2_x, plot_traj_f2_y, linewidth=0.8, color=traj_color)     # 跟随者2轨迹
                    plot_traj_x = []
                    plot_traj_y = []
                    plot_traj_f1_x = []
                    plot_traj_f1_y = []
                    plot_traj_f2_x = []
                    plot_traj_f2_y = []
        else:
            plt.plot(leader_x, leader_y, linewidth=2.0, color=traj_color, label="$trajectory1$")
            plt.plot(follower1_x, follower1_y, linewidth=0.8, color=traj_color)     # 跟随者1轨迹
            plt.plot(follower2_x, follower2_y, linewidth=0.8, color=traj_color)     # 跟随者2轨迹

        # 绘制编队队形
        for l1_x, l1_y, f1_x, f1_y, f2_x, f2_y in zip(leader_x, leader_y, follower1_x, follower1_y, follower2_x, follower2_y):
            if ((self.key1 % 50) == 0) or (self.key1 == len(leader_x) -1) :
                X = []
                Y = []
                X.append(l1_x)
                X.append(f1_x)
                X.append(f2_x)
                Y.append(l1_y)
                Y.append(f1_y)
                Y.append(f2_y)
                triangles = tri.Triangulation(X, Y)
                plt.triplot(triangles, color=formation_color, linestyle=(0, (5, 3)), linewidth=1.0)
            self.key1 = self.key1 + 1
    
    def plot_env(self):
        # 设置目标点和障碍物
        # 目标点
        # plt.legend(frameon=True, loc="upper right", fontsize='small')  # 分别为图例无边框、图例放在右上角、图例大小
        # x_major_locator = MultipleLocator(1)                    # 设置显示的间距
        # y_major_locator = MultipleLocator(1)
        ax = plt.gca()                                          # ax为两条坐标轴的实例
        # ax.xaxis.set_major_locator(x_major_locator)             # 把x轴的主刻度设置为1的倍数
        # ax.yaxis.set_major_locator(y_major_locator)             # 把y轴的主刻度设置为1的倍数
        plt.axis('equal')                                       # 把坐标轴单位长度都变的一样
        plt.xlim(-8.0, 8.0)  # 把x轴的刻度范围设置为-7.5到7.5，因为0.5不满一个刻度间隔，所以数字不会显示出来，但是能看到一点空白
        plt.ylim(-8.0, 7.0)
        ax.add_artist(Circle(xy=(-5.0, 5.0), radius=0.4, color='orangered'))
        ax.add_artist(Circle(xy=(0.0, 5.0), radius=0.4, color='orangered'))
        ax.add_artist(Circle(xy=(5.0, 5.0), radius=0.4, color='orangered'))

        # 障碍物
        ax.add_artist(Circle(xy=(1.0, -5), radius=0.2, color='g'))   # A
        ax.add_artist(Circle(xy=(-3, -6), radius=0.2, color='g'))  # B
        ax.add_artist(Circle(xy=(4, -3), radius=0.2, color='g'))   # C
        ax.add_artist(Circle(xy=(-6, -2), radius=0.2, color='g'))  # D
        ax.add_artist(Circle(xy=(7, -1), radius=0.2, color='g'))   # E
        ax.add_artist(Circle(xy=(0, -1), radius=0.2, color='g'))   # F
        ax.add_artist(Circle(xy=(-3, -1), radius=0.2, color='g'))  # G
        ax.add_artist(Circle(xy=(-8, 0), radius=0.2, color='g'))   # H
        ax.add_artist(Circle(xy=(4, 1), radius=0.2, color='g'))    # I
        ax.add_artist(Circle(xy=(-5, 1), radius=0.2, color='g'))   # J
        ax.add_artist(Circle(xy=(0, 3), radius=0.2, color='g'))    # K
        ax.add_artist(Circle(xy=(-3, 3), radius=0.2, color='g'))   # L
        ax.add_artist(Circle(xy=(-2, 1), radius=0.2, color='g'))   # X


if __name__ == '__main__':
    print(dirPath)
    with_auto_path = dirPath + "/ExperimentalData"
    without_auto_path = dirPath + "/ExperimentalData/yzg-brain_auto_A-1.csv"
    processor = Processor()
    processor.plot_trajs(with_auto_path, auto=True) # 把多个轨迹绘制到一个图中
    # processor.process_all(without_auto_path) # 包括计算指标，每个轨迹单独出图
    # processor.plot_traj(cur_file_path = "/home/chiweiming/catkin_ws/src/cwm_multi_nav/cwm_nav/script/ExperimentalData/YZG/withAuto/Hands-A-1.csv")
